from rest_framework.views import APIView
from rest_framework.response import Response
from rest_framework.renderers import JSONRenderer
from rest_framework.parsers import MultiPartParser
from django.core.files.storage import default_storage
from django.core.files.base import ContentFile
from django.conf import settings
import os
import cv2
from modelscope.outputs import OutputKeys
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
import time
import numpy as np
from rest_framework import status
import time

model_list = {
  '01': {
    'task': Tasks.image_color_enhancement,
    'model': 'damo/cv_csrnet_image-color-enhance-models',
  },
  '02': {
    'task': Tasks.image_portrait_stylization,
    'model': 'damo/cv_unet_person-image-cartoon-handdrawn_compound-models',
  },
  '03': {
    'task': Tasks.image_colorization,
    'model': 'damo/cv_ddcolor_image-colorization',
  },
  '04': {
    'task': Tasks.image_portrait_enhancement,
    'model': 'damo/cv_gpen_image-portrait-enhancement',
  },
  '05': {
    'task': Tasks.image_portrait_stylization, 
    'model': 'damo/cv_unet_person-image-cartoon-artstyle_compound-models',
  },
  '06': {
    'task': Tasks.image_portrait_stylization,
    'model': 'lskhh/moran-cv_unet_person-image-cartoon-genshin_compound-models',
  },
  '07': {
    'task': Tasks.image_portrait_stylization,
    'model': 'damo/cv_unet_person-image-cartoon-sd-design_compound-models',
  },
  '08': {
    'task': Tasks.image_portrait_stylization,
    'model': 'damo/cv_unet_person-image-cartoon-sd-illustration_compound-models',
  },
  '09': {
    'task': Tasks.image_portrait_stylization,
    'model': 'damo/cv_unet_person-image-cartoon-sketch_compound-models'
  },
}
class Img(APIView):
  renderer_classes = [JSONRenderer]
  parser_classes = [MultiPartParser]

  def post(self, request):
    file_obj = request.data['file']
    file_name = request.data['name']
    model_type = request.data['modelType']

    timestamp = time.time()

    file_name = f'{timestamp}_{file_name}'

    # 保存文件到MEDIA_ROOT指定的目录下
    relative_path = default_storage.save(
      os.path.join('uploads', file_name),
      ContentFile(file_obj.read())
    )
    absolute_path = os.path.join(settings.MEDIA_ROOT, relative_path)
    # img_path = '../..' + relative_path
    print('图片路径 ', absolute_path, model_list[model_type]['model'])
    time.sleep(2)
    img_cartoon = pipeline(
      model_list[model_type]['task'] or Tasks.image_portrait_stylization, 
      # model='damo/cv_unet_person-image-cartoon_compound-models'
      model= model_list[model_type]['model'] or 'damo/cv_unet_person-image-cartoon-3d_compound-models'
    )
    result = img_cartoon(absolute_path)
    # 确保输出是有效的图像数据
    output_img = result[OutputKeys.OUTPUT_IMG]
    output_path = ''
    if isinstance(output_img, np.ndarray):
      output_path = os.path.join(settings.MEDIA_ROOT, 'cartoon', file_name)
      os.makedirs(os.path.dirname(output_path), exist_ok=True)
      cv2.imwrite(output_path, output_img)
    else:
      return Response({'error': 'Invalid output image data.'}, status=status.HTTP_500_INTERNAL_SERVER_ERROR)

    return Response({
        'data': 'http://127.0.0.1:8000/media/' + 'cartoon/' + file_name,
        'msg': '上传文件成功，图片已转换并保存。',
        'code': status.HTTP_200_OK
      },
      status=status.HTTP_200_OK
    )
    # cv2.imwrite('result.png', result[OutputKeys.OUTPUT_IMG])
    # print('结束')
    # return Response('上传文件', content_type="application/json")
